#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
机器人TCP客户端模拟器
用于测试TCP服务器的机器人连接、注册和数据传输功能
"""

import socket
import json
import time
import threading
import random

class RobotSimulator:
    def __init__(self, server_host='localhost', server_port=8888, robot_sn=None):
        self.server_host = server_host
        self.server_port = server_port
        self.robot_sn = robot_sn or f"ROBOT_{random.randint(1000, 9999)}"
        self.socket = None
        self.running = False
        self.registered = False
        
    def connect(self):
        """连接到TCP服务器"""
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.connect((self.server_host, self.server_port))
            print(f"[{self.robot_sn}] 已连接到服务器 {self.server_host}:{self.server_port}")
            return True
        except Exception as e:
            print(f"[{self.robot_sn}] 连接失败: {e}")
            return False
    
    def disconnect(self):
        """断开连接"""
        self.running = False
        if self.socket:
            self.socket.close()
            print(f"[{self.robot_sn}] 已断开连接")
    
    def send_message(self, message):
        """发送JSON消息"""
        try:
            json_data = json.dumps(message, ensure_ascii=False)
            data = (json_data + '\n').encode('utf-8')
            self.socket.send(data)
            print(f"[{self.robot_sn}] 发送: {json_data}")
        except Exception as e:
            print(f"[{self.robot_sn}] 发送消息失败: {e}")
    
    def receive_messages(self):
        """接收消息线程"""
        buffer = b''
        while self.running:
            try:
                data = self.socket.recv(1024)
                if not data:
                    break
                
                buffer += data
                while b'\n' in buffer:
                    line, buffer = buffer.split(b'\n', 1)
                    if line:
                        try:
                            message = json.loads(line.decode('utf-8'))
                            print(f"[{self.robot_sn}] 收到: {message}")
                            self.handle_message(message)
                        except json.JSONDecodeError:
                            print(f"[{self.robot_sn}] JSON解析错误: {line}")
            except Exception as e:
                if self.running:
                    print(f"[{self.robot_sn}] 接收消息错误: {e}")
                break
    
    def handle_message(self, message):
        """处理收到的消息"""
        msg_type = message.get('type')
        if msg_type == 'register_ack':
            if message.get('status') == 'success':
                self.registered = True
                print(f"[{self.robot_sn}] 注册成功!")
            else:
                print(f"[{self.robot_sn}] 注册失败: {message.get('message')}")
        elif msg_type == 'heartbeat_ack':
            print(f"[{self.robot_sn}] 心跳响应收到")
    
    def register(self):
        """注册机器人"""
        register_msg = {
            "type": "register",
            "sn": self.robot_sn,
            "version": "1.0.0",
            "model": "TestRobot"
        }
        self.send_message(register_msg)
    
    def send_data(self):
        """发送机器人数据"""
        if not self.registered:
            return
            
        # 模拟机器人位置和温度数据
        data_msg = {
            "type": "data",
            "timestamp": int(time.time() * 1000),
            "position": {
                "x": round(random.uniform(-100, 100), 2),
                "y": round(random.uniform(-100, 100), 2),
                "z": round(random.uniform(0, 10), 2)
            },
            "temperature": round(random.uniform(20, 40), 1),
            "battery": random.randint(20, 100),
            "status": "working"
        }
        self.send_message(data_msg)
    
    def send_heartbeat(self):
        """发送心跳"""
        if not self.registered:
            return
            
        heartbeat_msg = {
            "type": "heartbeat",
            "timestamp": int(time.time() * 1000)
        }
        self.send_message(heartbeat_msg)
    
    def run_simulation(self):
        """运行模拟"""
        if not self.connect():
            return
        
        self.running = True
        
        # 启动接收消息线程
        receive_thread = threading.Thread(target=self.receive_messages)
        receive_thread.daemon = True
        receive_thread.start()
        
        # 等待连接稳定
        time.sleep(0.5)
        
        # 注册机器人
        self.register()
        
        # 等待注册成功
        timeout = 5
        while not self.registered and timeout > 0:
            time.sleep(0.1)
            timeout -= 0.1
        
        if not self.registered:
            print(f"[{self.robot_sn}] 注册超时，退出模拟")
            self.disconnect()
            return
        
        # 开始发送数据和心跳
        data_count = 0
        heartbeat_count = 0
        
        try:
            while self.running:
                # 每2秒发送一次数据
                if data_count % 20 == 0:  # 0.1 * 20 = 2秒
                    self.send_data()
                
                # 每10秒发送一次心跳
                if heartbeat_count % 100 == 0:  # 0.1 * 100 = 10秒
                    self.send_heartbeat()
                
                data_count += 1
                heartbeat_count += 1
                time.sleep(0.1)
                
        except KeyboardInterrupt:
            print(f"[{self.robot_sn}] 用户中断")
        finally:
            self.disconnect()

def main():
    print("机器人TCP客户端模拟器")
    print("用法: python robot_simulator.py")
    print("连接到默认服务器 localhost:8888")
    print("按 Ctrl+C 停止模拟")
    print()
    
    # 创建多个机器人进行测试
    robots = []
    robot_count = 2  # 可以调整机器人数量
    
    for i in range(robot_count):
        robot = RobotSimulator(robot_sn=f"ROBOT_{1001 + i}")
        robots.append(robot)
    
    # 启动机器人模拟线程
    threads = []
    for robot in robots:
        thread = threading.Thread(target=robot.run_simulation)
        thread.daemon = True
        threads.append(thread)
        thread.start()
        time.sleep(0.5)  # 错开连接时间
    
    try:
        # 等待所有线程
        for thread in threads:
            thread.join()
    except KeyboardInterrupt:
        print("\n正在停止所有机器人模拟...")
        for robot in robots:
            robot.disconnect()

if __name__ == "__main__":
    main()